
import torch
from eval import show_without_plt
import matplotlib.pyplot as plt
import numpy as np
# 中文显示
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False


def show_plot(load_tensors_1, load_tensors_2):
    counts_1, counts_2 = [], []
    bins = []
    # 遍历加载的张量
    for tensor1, tensor2 in zip(load_tensors_1, load_tensors_2):
        temp1 = show_without_plt(tensor1)
        temp2 = show_without_plt(tensor2)
        counts_1.append(temp1["counts"])
        counts_2.append(temp2["counts"])
        bins = temp1["bins"]  # 假设两组数据的 bins 是一致的

    # 计算平均值
    count1_avg = np.array(counts_1).mean(axis=0)
    count2_avg = np.array(counts_2).mean(axis=0)

    # 转换区间为字符串形式
    bin_labels = [f"{bins[i]}-{bins[i + 1]}" for i in range(len(bins) - 1)]
    print(bin_labels)

    # 绘制柱状图
    plt.figure(figsize=(16, 8))
    width = 0.4  # 每组柱子的宽度
    x = np.arange(len(bin_labels))  # x 轴的位置（以区间编号为中心）

    # 绘制两组柱子
    bars1 = plt.bar(x - width / 2, count1_avg, width=width, label="旋转前", color='skyblue', edgecolor='black')
    bars2 = plt.bar(x + width / 2, count2_avg, width=width, label="旋转后", color='salmon', edgecolor='black')
    title = "featuremap全局平均数值分布"
    plt.title(title, fontsize=15)

    # 设置横轴和纵轴标签
    plt.xlabel("区间", fontsize=12)
    plt.ylabel("统计数", fontsize=12)

    # 设置 x 轴刻度和标签
    plt.xticks(x, bin_labels, ha='center', fontsize=10, rotation=45)

    # 添加图例
    plt.legend(fontsize=12)

    # 添加网格
    plt.grid(axis='y', linestyle='--', linewidth=0.5)

    # 在柱状图上标注数值
    for bars in [bars1, bars2]:
        for bar in bars:
            height = bar.get_height()
            plt.text(bar.get_x() + bar.get_width() / 2, height,
                     f'{height:.0f}', ha='center', va='bottom', fontsize=10, color='black')

    # 显示图形
    plt.tight_layout()

    save_path = title
    if save_path:
        plt.savefig(save_path, dpi=300)  # 保存为高分辨率图片
        print(f"Figure saved to {save_path}")

    plt.show()


def show_mut_plot(load_tensors_1, load_tensors_2):
    counts_1, counts_2 = [], []
    bins = []

    # 遍历加载的张量
    for tensor1, tensor2 in zip(load_tensors_1, load_tensors_2):
        temp1 = show_without_plt(tensor1)
        temp2 = show_without_plt(tensor2)
        counts_1.append(temp1["counts"])
        counts_2.append(temp2["counts"])
        bins = temp1["bins"]

    # 转换区间为字符串形式
    bin_labels = [f"{bins[i]}-{bins[i + 1]}" for i in range(len(bins)-1)]
    print(bin_labels)

    fig, axs = plt.subplots(4, 8, figsize=(80, 25))  # 4行8列的子图
    width = 0.4  # 每组柱子的宽度

    for i, ax in enumerate(axs.flat):
        # 绘制两组柱子
        x = np.arange(len(bin_labels))
        bars1 = ax.bar(x - width / 2, counts_1[i], width=width, label="旋转前", color='skyblue', edgecolor='black')
        bars2 = ax.bar(x + width / 2, counts_2[i], width=width, label="旋转后", color='salmon', edgecolor='black')

        ax.set_title(f"第{i+1}个block的featuremap数值分布", fontsize=15)

        # 设置横轴和纵轴标签
        ax.set_xlabel(f"区间", fontsize=12)
        ax.set_ylabel("统计数", fontsize=12)

        # 设置 x 轴刻度和标签
        ax.set_xticks(x)
        ax.set_xticklabels(bin_labels, ha='center', fontsize=10 ,rotation=45)

        # 添加网格
        ax.grid(axis='y', linestyle='--', linewidth=0.5)

        # 在柱状图上标注数值
        for bars in [bars1, bars2]:
            for bar in bars:
                height = bar.get_height()
                ax.text(bar.get_x() + bar.get_width() / 2, height,
                        f'{height:.0f}', ha='center', va='bottom', fontsize=8, color='black')

        ax.legend(fontsize=12)
    plt.tight_layout()
    save_path = "各个block的featuremap数值分布"
    if save_path:
        plt.savefig(save_path, dpi=300)  # 保存为高分辨率图片
        print(f"Figure saved to {save_path}")
    plt.show()


if __name__ == '__main__':
    # 加载数据
    # load_tensors_1 = torch.stack([torch.load(f"../ffn/second_line_input/tensor{i}.pth", weights_only=True) for i in range(32)])
    # load_tensors_2 = torch.matmul(load_tensors_1, torch.load("../1-1-2.pth", weights_only=True).to("cpu"))

    load_tensors_1 = torch.load('original_input.pth', weights_only=True)
    load_tensors_2 = torch.load('modified_input.pth', weights_only=True)

    # 显示单个图和多子图
    show_plot(load_tensors_1, load_tensors_2)
    show_mut_plot(load_tensors_1, load_tensors_2)
